

import sys
import os
import json
import redis
import time
import shutil
from datetime import datetime
from PySide6.QtCore import Qt, QThread, Signal, Slot, QDateTime, QTimer, QObject
import requests

from gui.schemas.enumeration import AnnotationStatus, UserRole

def _parse_time(time_str):
    """解析时间字符串为datetime对象"""
    if not time_str:
        return None
    try:
        if "+" in time_str or "Z" in time_str:
            return datetime.fromisoformat(time_str.replace("Z", "+00:00"))
        else:
            return datetime.fromisoformat(time_str)
    except:
        return None
# 数据管理器 - 处理与Redis和本地文件的交互
class DataManager(QObject):
    """
    描述：数据管理器，处理与Redis和本地文件的交互
    1. 管理员可以创建工作，分配文件给用户，关闭工作，查看工作列表，查看文件列表，查看文件标注数据，保存文件标注数据，查看文件检查数据，保存文件检查数据，查看文件评论数据，保存文件评论数据，查看文件浏览数据，保存文件浏览数据，查看文件标注错误数据，保存文件标注错误数据，查看文件检查错误数据，保存文件检查错误数据，查看文件检查错误数据，保存文件检查错误数据，查看文件检查错误数据，保存文件检查错误数据，查看文件检查错误数据，保存文件检查错误数据，
    2. 标注员可以查看有自己的工作job、会获取最新的job和标注数据以及文件，只能修改属于自己的数据
    3. 检查员可以所有的工作job，只能修改评价标注的数据，标注目标就只能修改目标框评价
    4. 管理员直接操作redis，标注员和检查员需要下载到本地操作，
    Returns:
        _type_: _description_
    """
    connected = Signal(bool)
    
    def __init__(self, username="", role="", redis_host='localhost', redis_port=6379, redis_password="123456", redis_db=0):
        super().__init__()
        self.host = redis_host
        self.port = redis_port
        self.password = redis_password
        self.db = redis_db
        self.redis_client = None
        self.connected_status = False
        self.set_userinfo(username,role)
        self.models_key = "models"
        
    def set_userinfo(self,username,role):
        if username:
            self.username = username
            self.local_data_dir = os.path.join("./local_data", username)
            os.makedirs(self.local_data_dir, exist_ok=True)
        self.role = role
        
        
        
        
    def connect(self):
        """连接到Redis服务器"""
        try:
            self.redis_client = redis.Redis(
                host=self.host,
                port=self.port,
                password=self.password,
                db=self.db,
                decode_responses=True,
                socket_timeout=1,
            )
            # self.client.select(self.db)
            # 测试连接
            if self.redis_client.ping():
                self.connected_status = True
                self.connected.emit(True)
                # self.connection_checked.emit(True)
                return True
            else:
                self.connected_status = False
                self.connected.emit(False)
                # self.error_occurred.emit("无法ping通Redis服务器")
                # self.connection_checked.emit(False)
                return False
        except Exception as e:
            self.connected_status = False
            self.connected.emit(False)
            # self.error_occurred.emit(f"连接Redis出错: {str(e)}")
            # self.connection_checked.emit(False)
            return False

    def create_job(self, job_data,select_files):
        """创建新工作"""
        if UserRole.ADMIN in self.role and self.connected_status: # 具有管理员权限
            # job_id = datetime.now().isoformat().replace(":", "-")
            name = job_data["name"]
            # 获取模型信息
            model_name = job_data["model_name"]
            model_data = self.get_model(model_name) # 获取模型信息
            if model_data:
                job_data["labels"] = model_data["labels"]
                job_data["det_type"] = model_data["det_type"]
            job_data["status"] = "open"
            if "user_metadata" not in job_data:
                job_data["user_metadata"] = { # 用户自定义数据，不同用户不能修改
                    "config": {}, # 其他配置项，如快捷键等
                    "job": {
                        "current_id": 0 # 当前用户标注文件列表所在的索引
                    }
                }
            self.save_job(job_data)
            
            # 保存工作信息到本地, 管理员不需要保存工作信息到本地
            # job_dir = os.path.join(self.local_data_dir, job_data["name"])
            # os.makedirs(job_dir, exist_ok=True)
            # info_path = os.path.join(job_dir, ".info")
            # with open(info_path, "w") as f:
            #     json.dump(job_data, f, ensure_ascii=False, indent=2)
            
            # 搜索文件并添加到工作中
            base_url = job_data["base_url"]
            file_type = job_data["file_type"]
            file_ext = ".jpg" if file_type == "图片" else ".mp4"
            reviews = {}
            for filename in select_files:
                file_url = os.path.join(base_url, filename)
                if base_url.startswith("http"): # 可以是https
                    file_url = f"{base_url}/{filename}"
                    size = 0
                else:
                    size = os.path.getsize(file_url)
                file_data = {
                    "filename": filename,
                    "size": size,
                    "file_attributes" : {
                        "file_url": file_url, # 网络访问地址
                        "worker": "",
                        "transliterate": "", # OCR转写
                        "class_name": "", # 分类名
                    },
                    "regions": [],
                    "owner": "",
                    "status": "未分配",
                }
                reviews = {
                    "filename": filename,
                    "comment": "", # 网络访问地址
                    "regions": [],
                }
                self.save_file_annotation(name, file_data,reviews)
            return name
        else:
            raise PermissionError("只有管理员才能创建工作")
        
    def check_updatetime(self, redis_updatetime, locals_updatetime):
        return redis_updatetime > locals_updatetime
        
    def get_jobs(self):
        """获取所有工作,确保从最新源中获取数据即可"""
        jobs = {}
        close_jobs = []
        if self.connected_status: # 从redis获取
            job_keys = self.redis_client.keys("job:info:*")
            for job_key in job_keys:
                redis_job_info = self.redis_client.hgetall(job_key)
                # 读取 user_metadata 数据
                name = redis_job_info["name"]
                status = redis_job_info["status"]
                if status == "close":
                    close_jobs.append(name)
                    continue
                user_metadata = self.redis_client.hgetall(f"job:user_metadata:{name}:{self.username}")
                if user_metadata:
                    if "job" in user_metadata:
                        user_metadata["job"] =json.loads(user_metadata["job"])
                    if "config" in user_metadata:
                        user_metadata["config"] =json.loads(user_metadata["config"])
                    redis_job_info["user_metadata"] = user_metadata
                else:
                    redis_job_info["user_metadata"] = { # 用户自定义数据，不同用户不能修改
                        "config": {}, # 其他配置项，如快捷键等
                        "job": {
                            "current_id": 0 # 当前用户标注文件列表所在的索引
                        }
                    }
                
                # 将 JSON 字符串转换为列表
                if "labels" in redis_job_info:
                    redis_job_info["labels"] = json.loads(redis_job_info["labels"])
                redis_job_info["updatetime"] = _parse_time(redis_job_info["updatetime"])
                jobs[name] = redis_job_info
                
        if UserRole.ANNOTATOR in self.role or UserRole.EMEND in self.role: # 如果是 标注员或检查员，要从本地或redis中加载最新的数据
            job_keys = list(set(os.listdir(self.local_data_dir) + list(jobs.keys())))
            for job_key in job_keys:
                if job_key in close_jobs: # 对于关闭的任务, 确保全部数据上传后删除
                    continue
                # if job_key not in jobs: # 如果不存在，则从ti
                #     continue
                job_dir = os.path.join(self.local_data_dir, job_key)
                os.makedirs(job_dir, exist_ok=True)
                locals_job_info_file = os.path.join(self.local_data_dir, job_key, ".info")
                if os.path.exists(locals_job_info_file):
                    with open(locals_job_info_file, "r",encoding="utf-8") as f:
                        locals_job_info = json.load(f)
                        locals_job_info["updatetime"] = _parse_time(locals_job_info["updatetime"])
                        
                        if job_key in jobs: # 如果都存在就比较最新的
                            redis_job_info = jobs[job_key]
                            if self.check_updatetime(redis_job_info["updatetime"], locals_job_info["updatetime"]): # 如果redis数据比本地数据更新，则使用redis数据
                                self.save_job(redis_job_info)
                            else: # 如果本地数据新，则使用本地数据
                                jobs[job_key]["user_metadata"] = locals_job_info["user_metadata"] # 非管理员只能更新user_metadata数据
                                self.save_job(jobs[job_key])
                        else: # 如果redis数据不存在，则使用本地数据
                            jobs[job_key] = locals_job_info
                elif job_key in jobs: # 如果本地数据不存在，则使用redis数据,并保存到本地
                    redis_job_info = jobs[job_key]
                    self.save_job(redis_job_info)
        return jobs
    
    def save_job(self,job_info:dict):
        """保存job信息"""
        name = job_info["name"]
        job_info["createtime"] = datetime.now().isoformat()
        job_info["updatetime"] = datetime.now().isoformat()
        if UserRole.ANNOTATOR in self.role or UserRole.EMEND in self.role: # 如果是 标注员或检查员，要从本地或redis中加载最新的数据
            locals_job_info_file = os.path.join(self.local_data_dir, name, ".info")
            os.makedirs(os.path.dirname(locals_job_info_file),exist_ok=True)
            with open(locals_job_info_file, "w",encoding="utf-8") as f:
                job_info_str = json.dumps(job_info, ensure_ascii=False, indent=2)
                f.write(job_info_str)
        if self.connected_status: # 从redis获取
            user_metadata = job_info.pop("user_metadata")
            user_metadata["config"] = json.dumps(user_metadata.get("config",{}),ensure_ascii=False)
            user_metadata["job"] = json.dumps(user_metadata.get("job",{}),ensure_ascii=False)
            user_metadata["updatetime"] = datetime.now().isoformat()

            self.redis_client.hset(f"job:user_metadata:{name}:{self.username}", mapping=user_metadata)
            
            # 将列表类型的字段转换为 JSON 字符串
            if "labels" in job_info and isinstance(job_info["labels"], list):
                job_info["labels"] = json.dumps(job_info["labels"])
            self.redis_client.hset(f"job:info:{name}", mapping=job_info)

    def get_files_data(self, redis_data):
        file_datas = {}
        job_name = redis_data.get("name")
        filenames = []
        if self.connected_status: # 从redis获取
            job_filename_keys = self.redis_client.keys(f"job:files:{job_name}:*")
            for job_filename_key in job_filename_keys:
                filenames.append(job_filename_key.split(":")[-1])
            
        if UserRole.ANNOTATOR in self.role or UserRole.EMEND in self.role: # 如果是 标注员或检查员，要从本地或redis中加载最新的数据
            job_dir = os.path.join(self.local_data_dir, job_name)
            job_filenames = os.listdir(job_dir)
            for filename in job_filenames:
                basename, ext = os.path.splitext(filename)
                if filename not in filenames and ext in [".jpg"]:
                    filenames.append(filename)
                    
        for filename in filenames:
            file_annotation = self.get_file_annotation(redis_data,filename)
            if file_annotation:
                file_datas[file_annotation["filename"]] = file_annotation
        return file_datas
            
    def save_file_annotation(self, name, file_annotation = None,reviews = None):
        job_key = f"job:{name}"
        if file_annotation:
            file_annotation["updatetime"] = datetime.now().isoformat()
            filename = file_annotation["filename"]
        if reviews:
            reviews["updatetime"] = datetime.now().isoformat()
            filename = reviews["filename"]
        if UserRole.EMEND in self.role and reviews is not None: # 检测员，只能修改评论
            job_dir = os.path.join(self.local_data_dir, name)
            os.makedirs(job_dir, exist_ok=True)
            basename = os.path.splitext(filename)[0]
            file_data_path = os.path.join(job_dir, basename + ".reviews")
            with open(file_data_path, "w", encoding="utf-8") as f:
                json.dump(reviews, f, ensure_ascii=False, indent=2)
        if UserRole.ANNOTATOR in self.role and file_annotation is not None: # 标注员，不能修改评论
            job_dir = os.path.join(self.local_data_dir, name)
            os.makedirs(job_dir, exist_ok=True)
            basename = os.path.splitext(filename)[0]
            file_data_path = os.path.join(job_dir, basename + ".json")
            if "regions" in file_annotation and isinstance(file_annotation["regions"],str):
                file_annotation["regions"] = json.loads(file_annotation["regions"])
            if "file_attributes" in file_annotation and isinstance(file_annotation["file_attributes"],str):
                file_annotation["file_attributes"] = json.loads(file_annotation["file_attributes"])
            with open(file_data_path, "w", encoding="utf-8") as f:
                json.dump(file_annotation, f, ensure_ascii=False, indent=2)
                
        if UserRole.ADMIN in self.role or self.connected_status: # 具有管理员权限
            if file_annotation:
                file_annotation["regions"] = json.dumps(file_annotation["regions"],ensure_ascii=False)
                file_annotation["file_attributes"] = json.dumps(file_annotation["file_attributes"],ensure_ascii=False)
                if "reviews" in file_annotation:
                    file_annotation.pop("reviews")
                self.redis_client.hset(f"job:files:{name}:{filename}", mapping=file_annotation)
            if reviews:
                reviews["regions"] = json.dumps(reviews["regions"],ensure_ascii=False)
                self.redis_client.hset(f"job:reviews:{name}:{filename}", mapping=reviews)
        
    def get_file_annotation(self, redis_data, filename):
        """获取文件的标注数据"""
        job_name = redis_data.get("name")
        file_data = None
        redis_file_data = None
        redis_file_reviews = None
        if self.connected_status: # 从redis获取
            redis_file_data = self.redis_client.hgetall(f"job:files:{job_name}:{filename}")
            redis_file_reviews = self.redis_client.hgetall(f"job:reviews:{job_name}:{filename}")
            if redis_file_data:
                if "file_attributes" in redis_file_data:
                    redis_file_data["file_attributes"] = json.loads(redis_file_data["file_attributes"])
                if "regions" in redis_file_data:
                    redis_file_data["regions"] = json.loads(redis_file_data["regions"])
                if not redis_file_reviews:
                    redis_file_reviews = {
                        "comment": "", # 网络访问地址
                        "regions": [],
                    }
                else:
                    redis_file_reviews["regions"] = json.loads(redis_file_reviews["regions"])
                    
                file_data = redis_file_data
                if UserRole.EMEND in self.role or UserRole.ADMIN in self.role: # 检测员和管理员，获取所有数据
                    pass
                elif UserRole.ANNOTATOR in self.role: # 标注员，只能获取自己的数据
                    if redis_file_data["owner"] != self.username: # 不是自己的数据，跳过
                        return 
    

        if UserRole.ANNOTATOR in self.role or UserRole.EMEND in self.role: # 如果是 标注员或检查员，要从本地或redis中加载最新的数据
            job_dir = os.path.join(self.local_data_dir, job_name)
            locals_file_path = os.path.join(job_dir, filename)
            basename = os.path.splitext(filename)[0]
            locals_annon_file_path = os.path.join(job_dir, f"{basename}.json")
            locals_reviews_file_path = os.path.join(job_dir, f"{basename}.reviews")
            if os.path.exists(locals_annon_file_path):
                with open(locals_annon_file_path, "r", encoding="utf-8") as f: #判断标注文件
                    locals_annon_data = json.load(f)
                    if redis_file_data: # 如果都存在就比较最新的
                        file_data = redis_file_data
                        if not self.check_updatetime(redis_file_data["updatetime"], locals_annon_data["updatetime"]): # 如果redis数据比本地数据更新，则使用redis数据
                            file_data = locals_annon_data
                            self.save_file_annotation(job_name,file_data) # 保存文件数据
                    else: # 如果redis数据不存在，则使用本地数据
                        file_data = locals_annon_data
                        self.save_file_annotation(job_name,file_data) # 保存文件数据
            elif redis_file_data:
                self.save_file_annotation(job_name,redis_file_data) # 保存文件数据
                
            if not os.path.exists(locals_file_path):
                pass
                try:
                    file_url = file_data.get("file_attributes",{}).get("file_url")
                    if file_url.startswith("http"): # 可以是https
                        # file_url = f"{file_path}/{filename}"
                        r = requests.get(file_url)
                        with open(locals_file_path, "wb") as f:
                            f.write(r.content)
                    else:
                        # file_url = os.path.join(base_url, filename)
                        shutil.copy(file_url, locals_file_path)
                except Exception as e:
                    print(f"获取 {file_url} 文件失败: {e}")
                    traceback.print_exception(e)
                    return 
                
                
            if file_data is None: # 如果没有标注文件数据
                file_data = {
                    "filename": filename,
                    "size": os.path.getsize(locals_file_path),
                    "file_attributes" : {
                        # "file_url": file_url, # 网络访问地址
                        "worker": "",
                        "transliterate": "", # OCR转写
                        "class_name": "", # 分类名
                    },
                    "regions": [],
                    "owner": self.username,
                    "status": AnnotationStatus.ASSIGNED,
                }
                self.save_file_annotation(file_data) # 保存文件数据
                
            if os.path.exists(locals_reviews_file_path):
                with open(locals_reviews_file_path, "r", encoding="utf-8") as f:
                    locals_reviews_data = json.load(f)
                
                    if redis_file_data: # 如果都存在就比较最新的
                        if not self.check_updatetime(redis_file_reviews["updatetime"], locals_reviews_data["updatetime"]): # 如果redis数据比本地数据更新，则使用redis数据
                            self.save_file_annotation(job_name,reviews=locals_reviews_data) # 保存评论数据
                            file_data["reviews"] = locals_reviews_data
                    else: # 如果redis数据不存在，则使用本地数据
                        self.save_file_annotation(job_name,reviews=locals_reviews_data) # 保存评论数据
                        file_data["reviews"] = locals_reviews_data
            else:
                file_data["reviews"] = {
                    "filename": filename,
                    "comment": "", # 
                    "regions": [],
                }
                self.save_file_annotation(job_name,reviews=file_data["reviews"]) # 保存评论数据
                
        if "regions" in file_data and isinstance(file_data["regions"],str):
            file_data["regions"] = json.loads(file_data["regions"])
        if "file_attributes" in file_data and isinstance(file_data["file_attributes"],str):
            file_data["file_attributes"] = json.loads(file_data["file_attributes"])
        return file_data
    
        
    def assign_files(self, job_name, user_files):
        """分配文件给用户"""
        if UserRole.ADMIN in self.role: # 具有管理员权限
            for username, filenames in user_files.items():
                status = self.redis_client.hget(f"job:files:{job_name}:{filenames}","status")
                if status == AnnotationStatus.UNASSIGNED:
                    self.redis_client.hset(f"job:files:{job_name}:{filenames}", "owner", username)
                    self.redis_client.hset(f"job:files:{job_name}:{filenames}", "status", AnnotationStatus.ASSIGNED)
                    self.redis_client.hset(f"job:files:{job_name}:{filenames}", "updatetime", datetime.now().isoformat())
        
    
    def close_job(self, job_id):
        """关闭工作"""
        job_key = f"job:{job_id}"
        job_data = self.redis_client.hgetall(job_key)
        
        # 保存数据到后端数据库（接口调用，未实现）
        # save_to_database(job_data)
        
        # 删除Redis中的工作
        self.redis_client.delete(job_key)
        
        # 删除本地工作目录
        job_dir = os.path.join(self.local_data_dir, job_data["name"])
        if os.path.exists(job_dir):
            shutil.rmtree(job_dir)
    
    
    def add_model(self, model_name, labels, det_type):
        """添加模型"""
        model_data = {
            "labels": labels,
            "det_type": det_type
        }
        self.redis_client.hset(self.models_key, model_name, json.dumps(model_data))
    
    def get_all_models(self):
        """获取所有模型"""
        models = self.redis_client.hgetall(self.models_key)
        for key in list(models.keys()):
            models[key] = json.loads(models[key])
        return models
    
    def get_model(self, model_name):
        """获取单个模型信息"""
        model_data = self.redis_client.hget(self.models_key, model_name)
        if model_data:
            return json.loads(model_data)
        return None
